#!/usr/bin/env python3
from mininet.net import Mininet
from mininet.topo import Topo
from mininet.link import TCLink
from mininet.log import setLogLevel, info
from mininet.cli import CLI
from mininet.node import OVSSwitch  # 改为 OVSSwitch
import os
import time

class ECMPTopo(Topo):
    def build(self):
        # 创建6个交换机
        s1 = self.addSwitch('s1', cls=OVSSwitch)  # 指定使用 OVSSwitch
        s2 = self.addSwitch('s2', cls=OVSSwitch)
        s3 = self.addSwitch('s3', cls=OVSSwitch)
        s4 = self.addSwitch('s4', cls=OVSSwitch)
        s5 = self.addSwitch('s5', cls=OVSSwitch)
        s6 = self.addSwitch('s6', cls=OVSSwitch)

        # 创建2个主机
        h1 = self.addHost('h1', ip='10.0.1.1/24', mac='08:00:00:00:01:01')
        h2 = self.addHost('h2', ip='10.0.2.2/24', mac='08:00:00:00:02:02')

        # 添加链路
        self.addLink(h1, s1, port1=1, port2=1)
        self.addLink(s1, s2, port1=2, port2=1)
        self.addLink(s1, s3, port1=3, port2=1)
        self.addLink(s1, s4, port1=4, port2=1)
        self.addLink(s1, s5, port1=5, port2=1)
        self.addLink(s2, s6, port1=2, port2=2)
        self.addLink(s3, s6, port1=2, port2=3)
        self.addLink(s4, s6, port1=2, port2=4)
        self.addLink(s5, s6, port1=2, port2=5)
        self.addLink(s6, h2, port1=1, port2=1)

def run_topology():
    setLogLevel('info')

    # 创建目录
    os.makedirs('logs', exist_ok=True)
    os.makedirs('pcaps', exist_ok=True)

    # 创建拓扑
    topo = ECMPTopo()

    # 创建网络 - 使用 OVSSwitch
    net = Mininet(topo=topo,
                  switch=OVSSwitch,  # 使用 OVSSwitch
                  controller=None,
                  link=TCLink,
                  autoSetMacs=True)

    # 启动网络
    net.start()
    info('*** Network started ***\n')

    # 配置主机路由
    info('*** Configuring host routes...\n')

    # h1 配置
    h1 = net.get('h1')
    h1.cmd('route del default dev h1-eth0 2>/dev/null')
    h1.cmd('route add default gw 10.0.1.10 dev h1-eth0')
    h1.cmd('arp -i h1-eth0 -s 10.0.1.10 08:00:00:00:01:00')
    h1.cmd('route add -host 10.0.0.1 gw 10.0.1.10 dev h1-eth0')
    h1.cmd('route add -host 10.0.0.2 gw 10.0.1.10 dev h1-eth0')

    # h2 配置
    h2 = net.get('h2')
    h2.cmd('route del default dev h2-eth0 2>/dev/null')
    h2.cmd('route add default gw 10.0.2.20 dev h2-eth0')
    h2.cmd('arp -i h2-eth0 -s 10.0.2.20 08:00:00:00:02:00')
    h2.cmd('route add -host 10.0.0.1 gw 10.0.2.20 dev h2-eth0')
    h2.cmd('route add -host 10.0.0.2 gw 10.0.2.20 dev h2-eth0')

    info('*** Starting P4 switches...\n')

    # 启动 P4 交换机
    switches = ['s1', 's2', 's3', 's4', 's5', 's6']
    thrift_ports = [9091, 9092, 9093, 9094, 9095, 9096]

    for i, switch_name in enumerate(switches):
        switch = net.get(switch_name)
        thrift_port = thrift_ports[i]

        # 构建接口参数
        if switch_name == 's1':
            interface_args = '-i 1@s1-eth1 -i 2@s1-eth2 -i 3@s1-eth3 -i 4@s1-eth4 -i 5@s1-eth5'
        elif switch_name == 's2':
            interface_args = '-i 1@s2-eth1 -i 2@s2-eth2'
        elif switch_name == 's3':
            interface_args = '-i 1@s3-eth1 -i 2@s3-eth2'
        elif switch_name == 's4':
            interface_args = '-i 1@s4-eth1 -i 2@s4-eth2'
        elif switch_name == 's5':
            interface_args = '-i 1@s5-eth1 -i 2@s5-eth2'
        elif switch_name == 's6':
            interface_args = '-i 1@s6-eth1 -i 2@s6-eth2 -i 3@s6-eth3 -i 4@s6-eth4 -i 5@s6-eth5'

        # 启动 simple_switch_grpc
        cmd = f'simple_switch_grpc --thrift-port {thrift_port} --device-id {i} {interface_args} build/load_balance.json > logs/{switch_name}.log 2>&1 &'
        switch.cmd(cmd)
        info(f'*** Started {switch_name} on thrift port {thrift_port}\n')

    info('*** Waiting for switches to initialize...\n')
    time.sleep(10)

    info('*** Topology is running. Start controller with: make controller\n')
    info('*** Run tests from h1 and h2 terminals\n')

    # 进入 CLI
    CLI(net)

    # 清理
    net.stop()
    info('*** Network stopped ***\n')

if __name__ == '__main__':
    run_topology()
